// 2024/7/19: 首个版本
// 2024/7/22: 发送包中增加body.Length内容
// 2024/7/22: 新增RemoteSegmentRestore
// 2024/10/24: Restore解决可能导致的乱序问题，且移除所有Segment包
// 2025/7/23: 支持通讯协议v2
// 2025/7/25: 新增RemoteSegment.TryZip。RemoteSegmentRestore支持Zipped协议

using System;
using System.Collections.Generic;
using System.Text;
using System.Text.Json;
using System.IO;
using ICSharpCode.SharpZipLib.GZip;

namespace RemoteConsole
{
    class RemoteSegment
    {
        public static (bool, byte[]) TryZip(byte[] protoBinary)
        {
            if (protoBinary.Length >= MinZipSize && protoBinary.Length <= MaxZipSize)
            {
                using (var inputStream = new MemoryStream(protoBinary))
                {
                    using (var outputStream = new MemoryStream())
                    {
                        var header = new byte[4];
                        BitConverter.TryWriteBytes(header, (int)Protocol.Zipped);
                        outputStream.Write(header, 0, 4);
                        GZip.Compress(inputStream, outputStream, false);
                        if (outputStream.Length <= (1.0f - MinZipRatio) * protoBinary.Length)
                        {
                            return (true, outputStream.ToArray());
                        }
                        else return (false, protoBinary);
                    }
                }
            }
            else return (false, protoBinary);
        }

        public static byte[][] Generate(byte[] protoBinary, Protocol protocol)
        {
            if (protoBinary.Length <= MaxSegmentSize) return [protoBinary];

            int totalSegments = (int)Math.Ceiling((float)protoBinary.Length / SegmentSizeHint);
            if ((float)protoBinary.Length / (totalSegments - 1) <= MaxSegmentSize) totalSegments--;

            var segmentSize = (int)Math.Ceiling((float)protoBinary.Length / totalSegments);
            var guid = Guid.NewGuid().ToString().Replace("-", "");

            var segments = new byte[totalSegments][];
            var segmentHeadIndex = 0;
            for (int i = 0; i < totalSegments; i++)
            {
                var thisSegmentSize = Math.Min(segmentSize, protoBinary.Length - segmentHeadIndex);

                var proto = new SegmentProto
                {
                    GUID = guid,
                    Protocol = protocol,
                    TotalSize = protoBinary.Length,
                    TotalSegments = totalSegments,
                    SegmentIndex = i,
                    SegmentSize = thisSegmentSize,
                };

                var body = Encoding.UTF8.GetBytes(JsonSerializer.Serialize(proto));
                var binary = new byte[8 + body.Length + thisSegmentSize];
                BitConverter.TryWriteBytes(binary.AsSpan(0, 4), (int)Protocol.Segment);
                BitConverter.TryWriteBytes(binary.AsSpan(4, 4), body.Length);
                Array.Copy(body, 0, binary, 8, body.Length);
                Array.Copy(protoBinary, segmentHeadIndex, binary, 8 + body.Length, thisSegmentSize);

                segments[i] = binary;
                segmentHeadIndex += thisSegmentSize;
            }

            return segments;
        }

        private const int SegmentSizeHint = 16384; // 16k
        private const float Tolerance = 10; // %
        private const int MaxSegmentSize = (int)(SegmentSizeHint * (1.0f + 0.01f * Tolerance));
        private const int MinZipSize = 200; // bytes
        private const int MaxZipSize = 1000000; // bytes
        private const float MinZipRatio = 0.1f;
    }

    class RemoteSegmentRestore
    {
        public void Reset()
        {
            ctxs.Clear();
            lastCheckTimeout = DateTime.Now;
        }

        public void Restore(List<byte[]> receivedPackets)
        {
            var restoredPackets = new List<byte[]>();
            foreach (var binary in receivedPackets)
            {
                if (binary.Length < 8)
                {
                    restoredPackets.Add(binary);
                    continue;
                }

                var protocol = (Protocol)BitConverter.ToInt32(binary.AsSpan(0, 4));
                var jsonSize = BitConverter.ToInt32(binary.AsSpan(4, 4));
                if (protocol != Protocol.Segment)
                {
                    restoredPackets.Add(tryUnzip(binary));
                    continue;
                }

                var proto = JsonSerializer.Deserialize<SegmentProto>(binary.AsSpan(8, jsonSize));
                if (binary.Length != 8 + jsonSize + proto.SegmentSize) continue;

                if (!ctxs.ContainsKey(proto.GUID))
                {
                    if (proto.SegmentIndex != 0) continue;
                    ctxs[proto.GUID] = new SegmentRestoreContext(proto.TotalSize, proto.Protocol);
                }
                var ctx = ctxs[proto.GUID];

                if (proto.SegmentIndex != ctx.NextIndex || proto.Protocol != ctx.Protocol || ctx.WrittenBytes + proto.SegmentSize > ctx.Buffer.Length)
                {
                    ctxs.Remove(proto.GUID);
                    continue;
                }

                Array.Copy(binary, 8 + jsonSize, ctx.Buffer, ctx.WrittenBytes, proto.SegmentSize);
                ctx.WrittenBytes += proto.SegmentSize;
                ctx.NextIndex++;
                ctx.LastReceivedTime = DateTime.Now;

                if (ctx.NextIndex == proto.TotalSegments)
                {
                    if (ctx.WrittenBytes == ctx.Buffer.Length && ctx.Buffer.Length >= 4 && 
                        (Protocol)BitConverter.ToInt32(ctx.Buffer.AsSpan(0, 4)) == ctx.Protocol)
                    {
                        restoredPackets.Add(tryUnzip(ctx.Buffer));
                    }
                    ctxs.Remove(proto.GUID);
                }
            }

            receivedPackets.Clear();
            receivedPackets.AddRange(restoredPackets);

            var toRemove = new List<String>();
            var now = DateTime.Now;
            if ((now - lastCheckTimeout).TotalSeconds > 4)
            {
                foreach (var pair in ctxs)
                {
                    var ctx = pair.Value;
                    if ((now - ctx.LastReceivedTime).TotalSeconds > 4) toRemove.Add(pair.Key);
                }
            }
            foreach (var key in toRemove) ctxs.Remove(key);
        }

        private byte[] tryUnzip(byte[] packet)
        {
            if (BitConverter.ToInt32(packet.AsSpan(0, 4)) != (int)Protocol.Zipped) return packet;

            using (var inputStream = new MemoryStream(packet, 4, packet.Length - 4))
            {
                using (var outputStream = new MemoryStream())
                {
                    GZip.Decompress(inputStream, outputStream, false);
                    return outputStream.ToArray();
                }
            }
        }

        private class SegmentRestoreContext
        {
            public byte[] Buffer { get; set; }
            public int WrittenBytes { get; set; }
            public int NextIndex { get; set; }
            public Protocol Protocol { get; set; }
            public DateTime LastReceivedTime { get; set; }

            public SegmentRestoreContext(int bufferSize, Protocol protocol)
            {
                Buffer = new byte[bufferSize];
                Protocol = protocol;
            }
        }

        private Dictionary<String, SegmentRestoreContext> ctxs = new Dictionary<string, SegmentRestoreContext>();
        private DateTime lastCheckTimeout = DateTime.Now;
    }
}